import albumentations as A
from albumentations.pytorch import ToTensorV2
from const import BATCH_SIZE

from dataset_utils import ICBHI_Dataset, Dataset
from utils.data import collate_fn
from torch.utils.data import DataLoader

# region transforms
def get_train_transform():
    return A.Compose([
        A.Flip(0.5),
        A.RandomRotate90(0.5),
        A.MotionBlur(p=0.2),
        A.MedianBlur(blur_limit=3, p=0.1),
        A.Blur(blur_limit=3, p=0.1),
        ToTensorV2(p=1.0),
    ], bbox_params={
        'format': 'pascal_voc',
        'label_fields': ['labels']
    })
# define the validation transforms
def get_valid_transform():
    return A.Compose([
        ToTensorV2(p=1.0),
    ], bbox_params={
        'format': 'pascal_voc', 
        'label_fields': ['labels']
    })
# endregion

train_dict, test_dict = ICBHI_Dataset().DatasetDivision(test_fold = 4)
# TrainDataset = Dataset(train_dict, get_train_transform())
# TestDataset = Dataset(test_dict, get_valid_transform())
TrainDataset = Dataset(train_dict)
TestDataset = Dataset(test_dict)

TrainLoader = DataLoader(
    TrainDataset,
    batch_size = BATCH_SIZE,
    shuffle=True,
    num_workers=0,
    collate_fn=collate_fn
)

TestLoader = DataLoader(
    TestDataset,
    batch_size = BATCH_SIZE,
    shuffle = False,
    num_workers = 0,
    collate_fn=collate_fn
)

print(f"Number of training samples: {len(TrainDataset)}")
print(f"Number of validation samples: {len(TestDataset)}\n")

